import matplotlib.pyplot as plt
plt.rcParams['font.sans-serif'] = ['SimHei']  # 设置中文字体
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from pymoo.algorithms.moo.nsga2 import NSGA2
from pymoo.optimize import minimize
from pymoo.core.problem import Problem
from pymoo.operators.crossover.sbx import SBX
from pymoo.operators.mutation.pm import PM
from pymoo.operators.sampling.rnd import FloatRandomSampling
from pymoo.visualization.scatter import Scatter


# ====================== 1. 数据准备 ======================
def prepare_data():
    """模拟生成上海市各区康养资源配置数据"""
    regions = ["浦东新区", "黄浦区", "静安区", "徐汇区", "长宁区", "普陀区", "虹口区", "杨浦区",
               "闵行区", "宝山区", "嘉定区", "金山区", "松江区", "青浦区", "奉贤区", "崇明区"]
    np.random.seed(42)

    # 老年人口数据（万人）
    elder_pop = np.random.randint(5, 35, size=len(regions))

    # 设施类型及参数
    facilities = {
        "三级医院": {"cost": 10000, "capacity": 3000, "radius": 8.0},
        "养老院": {"cost": 5000, "capacity": 500, "radius": 5.0},
        "社区康养中心": {"cost": 1000, "capacity": 200, "radius": 2.0}
    }

    # 生成距离矩阵（km）
    dist_matrix = np.random.uniform(1, 15, (len(regions), len(regions)))
    np.fill_diagonal(dist_matrix, 0)

    return regions, elder_pop, facilities, dist_matrix


# ====================== 2. 优化模型构建 ======================
class CareResourceAllocation(Problem):
    def __init__(self, regions, pop, facilities, dist_matrix, budget=1e5):
        self.regions = regions
        self.pop = pop
        self.facilities = facilities
        self.dist_matrix = dist_matrix
        self.fac_types = list(facilities.keys())

        n_regions = len(regions)
        n_facilities = len(facilities)

        # 变量：x_ij（是否建设）+ y_ij（服务量）
        super().__init__(
            n_var=n_regions * n_facilities * 2,
            n_obj=2,  # 目标1：最大化覆盖率，目标2：最小化成本
            n_constr=n_regions + 1,  # 需求约束+预算约束
            xl=0,
            xu=1
        )

        # 设施参数矩阵化
        self.cost_vec = np.array([f["cost"] for f in facilities.values()])
        self.cap_vec = np.array([f["capacity"] for f in facilities.values()])
        self.rad_vec = np.array([f["radius"] for f in facilities.values()])
        self.budget = budget

    def _evaluate(self, X, out, *args, **kwargs):
        n_individuals = X.shape[0]
        n_regions = len(self.regions)
        n_facilities = len(self.facilities)

        coverage = np.zeros(n_individuals)
        cost = np.zeros(n_individuals)
        g1 = np.zeros((n_individuals, n_regions))  # 需求约束
        g2 = np.zeros(n_individuals)  # 预算约束

        for k in range(n_individuals):
            # 解码决策变量
            x = X[k, :n_regions * n_facilities].reshape(n_regions, n_facilities)
            y = X[k, n_regions * n_facilities:].reshape(n_regions, n_facilities) * self.cap_vec

            # 计算覆盖率（考虑服务半径）
            total_coverage = 0
            for i in range(n_regions):
                for j in range(n_facilities):
                    if x[i, j] > 0.5:  # 建设阈值
                        # 计算该设施能服务的区域
                        served_regions = np.where(self.dist_matrix[i, :] <= self.rad_vec[j])[0]
                        served_pop = np.sum(self.pop[served_regions])
                        # 覆盖率=服务人口×利用率×距离衰减
                        utilization = min(1, y[i, j] / self.cap_vec[j])
                        total_coverage += served_pop * utilization * np.exp(
                            -self.dist_matrix[i, served_regions] / self.rad_vec[j]).mean()

            coverage[k] = -total_coverage  # 转为最小化

            # 计算总成本
            construction_cost = np.sum(x * self.cost_vec)
            operation_cost = 0.2 * np.sum(y * self.cost_vec / self.cap_vec)
            cost[k] = construction_cost + operation_cost

            # 需求约束：每个区域服务总量≥需求的10%
            g1[k, :] = 0.1 * self.pop - np.sum(y, axis=1)

            # 预算约束
            g2[k] = cost[k] - self.budget

        out["F"] = np.column_stack([coverage, cost])
        out["G"] = np.column_stack([g1, g2])


# ====================== 3. 模型求解 ======================
def solve_model(problem):
    algorithm = NSGA2(
        pop_size=50,
        sampling=FloatRandomSampling(),
        crossover=SBX(prob=0.9, eta=15),
        mutation=PM(eta=20),
        eliminate_duplicates=True
    )

    res = minimize(
        problem,
        algorithm,
        ('n_gen', 100),
        seed=42,
        verbose=True
    )
    return res


# ====================== 4. 结果分析 ======================
def analyze_results(res, problem):
    # 提取帕累托前沿
    front = res.F
    solutions = res.X

    # 筛选可行解
    feasible = np.all(res.G <= 0, axis=1)
    feasible_F = front[feasible]
    feasible_X = solutions[feasible]

    if len(feasible_F) == 0:
        print("警告：未找到严格可行解，展示最近似解")
        violation = np.sum(np.where(res.G > 0, res.G, 0), axis=1)
        best_idx = np.argmin(violation)
        best_solution = solutions[best_idx]
    else:
        # 选择覆盖率最高且成本合理的解
        norm_F = (feasible_F - feasible_F.min(axis=0)) / (feasible_F.max(axis=0) - feasible_F.min(axis=0))
        best_idx = np.argmin(norm_F[:, 0] + 0.3 * norm_F[:, 1])  # 权重偏好
        best_solution = feasible_X[best_idx]

    # 解码最优解
    n_regions = len(problem.regions)
    n_facilities = len(problem.facilities)
    x_ij = best_solution[:n_regions * n_facilities].reshape(n_regions, n_facilities)
    y_ij = best_solution[n_regions * n_facilities:].reshape(n_regions, n_facilities) * problem.cap_vec

    # 输出建设方案
    print("\n=== 最优资源配置方案 ===")
    construction_plan = []
    for i in range(n_regions):
        for j in range(n_facilities):
            if x_ij[i, j] > 0.5:
                construction_plan.append({
                    "区域": problem.regions[i],
                    "设施类型": problem.fac_types[j],
                    "建设规模": f"{y_ij[i, j]:.0f}/{problem.cap_vec[j]}",
                    "服务半径(km)": problem.rad_vec[j],
                    "覆盖区域": ", ".join(
                        np.array(problem.regions)[np.where(problem.dist_matrix[i, :] <= problem.rad_vec[j])[0]])
                })

    df_plan = pd.DataFrame(construction_plan)
    print(df_plan.to_string(index=False))

    # 计算绩效指标
    coverage = -front[best_idx, 0]
    total_cost = front[best_idx, 1]
    print(f"\n● 总覆盖率: {coverage:.1f}万老人（占总老年人口{coverage / np.sum(problem.pop) * 100:.1f}%）")
    print(f"● 总成本: {total_cost / 1e4:.2f}亿元（预算利用率{total_cost / problem.budget * 100:.1f}%）")

    # 可视化
    plt.figure(figsize=(12, 5))

    # 帕累托前沿
    plt.subplot(1, 2, 1)
    plt.scatter(-front[:, 0], front[:, 1], c='blue', alpha=0.3)
    if np.any(feasible):
        plt.scatter(-feasible_F[:, 0], feasible_F[:, 1], c='red', label='可行解')
    plt.scatter(-front[best_idx, 0], front[best_idx, 1], c='green', s=200, marker='*', label='最优解')
    plt.title("帕累托前沿")
    plt.xlabel("覆盖率（万人）")
    plt.ylabel("成本（万元）")
    plt.legend()

    # 区域资源配置
    plt.subplot(1, 2, 2)
    region_allocation = np.sum(x_ij > 0.5, axis=1)
    plt.barh(problem.regions, region_allocation)
    plt.title("各区设施建设数量")
    plt.tight_layout()
    plt.show()

    return df_plan


# ====================== 5. 主程序 ======================
if __name__ == "__main__":
    print("====== 康养资源优化配置模型 ======")

    # 数据准备
    regions, pop, facilities, dist_matrix = prepare_data()
    print(f"\n■ 模拟数据概况：")
    print(f"- 区域数量: {len(regions)}")
    print(f"- 总老年人口: {np.sum(pop)}万人")
    print(f"- 设施类型: {list(facilities.keys())}")

    # 模型构建
    problem = CareResourceAllocation(
        regions=regions,
        pop=pop,
        facilities=facilities,
        dist_matrix=dist_matrix,
        budget=8e4  # 8亿元预算
    )

    # 模型求解
    print("\n■ 开始优化求解...")
    res = solve_model(problem)

    # 结果分析
    print("\n■ 优化结果分析")
    df_plan = analyze_results(res, problem)

    # 保存结果
    df_plan.to_csv("康养资源优化方案.csv", index=False)
    print("\n结果已保存到：康养资源优化方案.csv")